# -*- coding: utf-8 -*-
"""
   File Name：     inference
   Description :  frcnn预测
   Author :       mick.yi
   date：          2019/2/13
"""
import argparse
import os
import sys

import matplotlib
import numpy as np
import tensorflow as tf
from tensorflow.python.framework import graph_util


matplotlib.use('Agg')

from faster_rcnn.preprocess.input import VocDataset
from faster_rcnn.utils import image as image_utils
from faster_rcnn.config import current_config as config
from pb_inference import pb_inference


def pb_model_load(model_path="./models/predict_pb/saved_model.pb"):
    print("model_exist: {}".format(os.path.exists(model_path)))
    output_graph_path = os.path.dirname(model_path)
    sess = tf.Session(graph=tf.Graph())

    tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.SERVING], output_graph_path)
    graph = tf.get_default_graph()

    tensor_names = [[v.name for v in op.values()] for op in sess.graph.get_operations()]
    print("tensor_names: {}".format(tensor_names))
    # print("operations: {}".format(sess.graph.get_operations()))
    '''
    In my case, I named my input and output tensors as
    input:0 and output:0 respectively
    '''
    return sess
    # y_pred = sess.run('output:0', feed_dict={'input:0': X_test})


def class_map_to_id_map(class_mapping):
    id_map = {}
    for k, v in class_mapping.items():
        id_map[v] = k
    return id_map



def main(args):
    os.environ["CUDA_VISIBLE_DEVICES"] = config.INFERENCE_GPU_ID
    # 覆盖参数
    config.IMAGES_PER_GPU = 1
    config.GPU_COUNT = 1
    # 加载数据
    dataset = VocDataset(config.voc_path, class_mapping=config.CLASS_MAPPING)
    dataset.prepare()
    all_img_info = [info for info in dataset.get_image_info_list() if info['type'] == 'test']  # 测试集

    # 加载模型
    # m = models.frcnn(config, stage='test')
    # m.load_weights(args.weight_path, by_name=True)
    # m.summary()

    # class map 转为 id map
    id_mapping = class_map_to_id_map(config.CLASS_MAPPING)

    def _show_inference(id, ax=None):
        image = image_utils.load_image(all_img_info[id]['filepath'])
        image, image_meta, _ = image_utils.resize_image_and_gt(image,
                                                               config.IMAGE_MAX_DIM,
                                                               all_img_info[id]['boxes'])
        # print("image_meta: {}".format(type(image_meta)))

        # image = image.tolist()
        # image_meta = image_meta.tolist()
        print("image.shape: {}".format(image.shape))
        print("image_meta: {}".format(image_meta.shape))
        # result = pb_inference.pb_http_request({"input_image": np.expand_dims(image, axis=0).tolist(),
        #                                        "input_image_meta": np.expand_dims(image_meta, axis=0).tolist()})
        input_gt_class_ids = np.asarray([[1, 1] for i in range(50)])
        input_gt_boxes = np.asarray(([[1, 1, 1, 1, 1] for i in range(50)]))
        sess = pb_model_load()
        # result = sess.run(["rcnn_class_loss:0", "rcnn_bbox_loss/Mean:0", "identity_gt_num/Identity:0",
        #                    "rpn_class_loss/Mean:0", "rpn_bbox_loss/Mean:0"],
        #                   feed_dict={"input_image:0": image,
        #                              "input_image_meta:0": image_meta,
        #                              "input_gt_class_ids:0": input_gt_class_ids,
        #                              "input_gt_boxes:0": input_gt_boxes})

        result = sess.run(["lambda_4/Identity_1:0", "lambda_4/Identity:0"],
                          feed_dict={"input_image:0": np.expand_dims(image, axis=0),
                                     "input_image_meta:0": np.expand_dims(image_meta, axis=0)})

        # result = pb_inference.pb_http_request({"input_image": image.tolist(),
        #                                        "input_image_meta": image_meta.tolist()})
        print("result: {}".format(result))
        # boxes, scores, class_ids, _, image_metas = m.predict(
        #     [np.expand_dims(image, axis=0), np.expand_dims(image_meta, axis=0)])
        # boxes = np_utils.remove_pad(boxes[0])
        # scores = np_utils.remove_pad(scores[0])[:, 0]
        # class_ids = np_utils.remove_pad(class_ids[0])[:, 0]


    # 随机展示9张图像
    image_ids = np.random.choice(len(all_img_info), 1, replace=False)
    for idx, image_id in enumerate(image_ids):
        _show_inference(image_id, None)
    # fig.savefig('demo_images/inferece_examples.{}.png'.format(np.random.randint(10)))


if __name__ == '__main__':
    parse = argparse.ArgumentParser()
    parse.add_argument("--weight_path", type=str, default=None, help="weight path")
    argments = parse.parse_args(sys.argv[1:])
    main(argments)
